package org.example.petitionplatformsystem.service.Impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import io.reactivex.Flowable;
import org.example.petitionplatformsystem.service.AiService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class AiServiceImpl implements AiService {
    @Autowired
    private Generation generation;
    @Value("${ai.api.key}")
    private String appKey;

    @Value("${ai.model}")
    private String model;

    // 使用ConcurrentHashMap来存储每个用户的对话历史
    private Map<Long, List<Message>> userConversationHistory = new ConcurrentHashMap<>();
    private Map<Long, Integer> userLastSentIndex = new ConcurrentHashMap<>();

    @Override
    public Flux<ServerSentEvent<String>> aiTalk(Long userId, String question, HttpServletResponse response) throws NoApiKeyException, InputRequiredException {
        // 获取或创建用户的对话历史
        List<Message> conversationHistory = userConversationHistory.computeIfAbsent(userId, k -> new ArrayList<>());
        final int[] lastSentIndex = {userLastSentIndex.getOrDefault(userId, 0)};

        // 添加当前问题到对话历史
        conversationHistory.add(Message.builder().role(Role.USER.getValue()).content(question).build());

        // 确保对话历史不超过6条
        if (conversationHistory.size() > 6) {
            conversationHistory = conversationHistory.subList(conversationHistory.size() - 6, conversationHistory.size());
        }

        // 构建请求参数
        GenerationParam generationParam = GenerationParam.builder()
                .model(model)
                .messages(conversationHistory)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .apiKey(appKey)
                .build();

        // 执行文本生成操作，并流式返回结果
        Flowable<GenerationResult> result = generation.streamCall(generationParam);

        List<Message> finalConversationHistory = conversationHistory;
        return Flux.from(result)
                .flatMap(m -> {
                    String newContent = m.getOutput().getChoices().get(0).getMessage().getContent();
                    // 提取增量部分
                    String incrementalContent = newContent.substring(lastSentIndex[0]);
                    lastSentIndex[0] = newContent.length();
                    userLastSentIndex.put(userId, lastSentIndex[0]);

                    System.out.println("Incremental Content: " + incrementalContent); // 打印增量内容

                    // 添加生成的回复到对话历史
                    finalConversationHistory.add(Message.builder().role(Role.ASSISTANT.getValue()).content(incrementalContent).build());
                    return Mono.just(ServerSentEvent.<String>builder().data(incrementalContent).build());
                })
                .doOnComplete(() -> {
                    // 重置索引
                    lastSentIndex[0] = 0;
                    userLastSentIndex.remove(userId); // 从map中删除用户的索引记录
                })
                .doOnNext(event -> {
                    System.out.println("Event Data: " + event.data()); // 打印事件数据
                })
                .onErrorResume(e -> {
                    // 错误处理
                    Map<String, Object> map = new HashMap<>();
                    map.put("code", "400");
                    map.put("message", "has mistake " + e.getMessage());
                    return Mono.just(ServerSentEvent.<String>builder().data(JsonUtils.toJson(map)).build());
                })
                .publishOn(Schedulers.boundedElastic());
    }
}